import rclpy
from rclpy.node import Node
from my_vision_interfaces.msg import TargetPose
from geometry_msgs.msg import TransformStamped
from tf2_ros import TransformBroadcaster
import numpy as np
import math

# (复用rvec到四元数的转换函数)
def rvec_to_quaternion(rvec):
    angle = np.linalg.norm(rvec)
    if angle < 1e-9: return [0.0, 0.0, 0.0, 1.0]
    axis = rvec.flatten() / angle
    half_angle = angle / 2; sin_half_angle = math.sin(half_angle)
    return [axis[0]*sin_half_angle, axis[1]*sin_half_angle, axis[2]*sin_half_angle, math.cos(half_angle)]

class PoseToTfNode(Node):
    def __init__(self):
        super().__init__('pose_to_tf_broadcaster')
        # 初始化TF广播器
        self.tf_broadcaster = TransformBroadcaster(self)
        # 订阅我们自己的位姿话题
        self.subscription = self.create_subscription(
            TargetPose,
            '/target_pose',
            self.pose_callback,
            10)
        self.get_logger().info('TF广播节点已启动，等待位姿数据...')

    def pose_callback(self, msg):
        t = TransformStamped()

        # --- 填充TF消息 ---
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = 'camera_link'  # 父坐标系：相机
        t.child_frame_id = 'target_frame' # 子坐标系：目标物体

        # 平移向量 (tvec)
        t.transform.translation.x = msg.tvec.x
        t.transform.translation.y = msg.tvec.y
        t.transform.translation.z = msg.tvec.z

        # 旋转向量 (rvec) 需要转换为四元数
        rvec = np.array([msg.rvec.x, msg.rvec.y, msg.rvec.z])
        qx, qy, qz, qw = rvec_to_quaternion(rvec)
        t.transform.rotation.x = qx
        t.transform.rotation.y = qy
        t.transform.rotation.z = qz
        t.transform.rotation.w = qw

        # 广播TF变换
        self.tf_broadcaster.sendTransform(t)

def main(args=None):
    rclpy.init(args=args)
    node = PoseToTfNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
